import cv2
import numpy as np
import torch

from torch.utils.data import Dataset
import os
from torch.nn.functional import one_hot
from PIL import Image


class YellowDataset(Dataset):
    def __init__(self, root,is_train):
        self.dataset = []
        dir = 'train' if is_train else 'test'
        sub_dir = os.path.join(root,dir)
        for filename in os.listdir(sub_dir):
            img_dir = os.path.join(sub_dir, filename)
            #print(img_dir)
            self.dataset.append(img_dir)

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, index):
        data = self.dataset[index]
        #HWC
        img = cv2.imread(data)
        #print(img.shape)
        img = img / 255
        #CHW
        img = torch.tensor(img).permute(2, 0, 1)
        #print(img.shape)
        #归一化
        data_list = data.split('.')
        label=int(data_list[1])
        position = data_list[2:6]
        position = [int(i) / 300 for i in position]
        sort = int(data_list[6])-1

        return np.float32(img),np.float32(label), np.float32(position), sort


if __name__ == '__main__':
    data_dir = 'D:\BaiduNetdiskWorkspace\homework\cv\pytorch-Single-target-detection-Minions\yellow_demo\yellow_demo\data'
    data = YellowDataset(data_dir,True)
    for i in data:
        print(i)
